#!/usr/bin/env python 
# -*- coding: utf-8 -*-
# @Time    : 2023/4/9 17:02
# @Author  : WJQ
# @Site    : 
# @File    : driver_wx.py
# @Software: PyCharm
import os, time, json
import traceback
from appium.options.android import UiAutomator2Options

from appium import webdriver
from appium.webdriver.common.touch_action import TouchAction
from appium.webdriver.common.appiumby import AppiumBy
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC

import conf
from browser.driver import DriverHandler
from browser.page import PageHandler
from utils.common import CrawlerLogger, open_app
from utils.get_app_json import copy_Pkg_Down_wx, unpack_and_copy_json_wx


class DriverHandlerWx(DriverHandler):
    def __init__(self, applets_name):
        super().__init__(applets_name, 86)
        self.user_hash = "8ad392dbce9a5722e419bfe9227bb819"

    @classmethod
    def create(cls, applets_name):
        self = DriverHandlerWx(applets_name)
        return self

    def get_wechat_url ( self ) -> str :
        # e.g. accumulationFund/pages/housFund/housFund.html?type=%E5%85%AC%E7%A7%AF%E9%87%91%E6%9F%A5%E8%AF%A2&tag=home
        # 确保已经附加到应用
        # self.attach_to_app('微信', 'wechat')
        # 从 Frida 脚本获取路径
        # url = self.script.exports.get_url()  # 这里应该返回一个字符串
        # print("get_wechat_url")
        url = self.driver.execute_script('return window.__route__+".html" + (window.__queryString__ ? "?"+window.__queryString__ : "");')
        # print(url)
        return url

    def get_wechat_appid ( self ) -> str :
        # e.g. wxbebb3cdd9b331046
        # 确保已经附加到应用
        # self.attach_to_app('微信', 'wechat')
        # 从 Frida 脚本获取路径
        appid = self.driver.execute_script("return window.__appId__;")
        # appid = self.script.exports.get_appid()  # 这里应该返回一个字符串
        return appid
    
    def startMiniApp(self, appid):
        # e.g. wxbebb3cdd9b331046
        # 确保已经附加到应用
        self.attach_to_app('微信', 'wechat')
        # rpc
        self.script.exports.start_mini_app(appid)

    def check_jump( self ) -> None:
        all_functions = self.driver.execute_script(
            "return Object.keys(wx).filter(key => typeof wx[key] === 'function');")
        if "reLaunch" not in all_functions :
            CrawlerLogger.error("[-] Error finding webview context, maybe is a game")
            exit()

    def find_top_window(self):
        '''
        切换成当前driver的window handle成最顶层的window
        '''
        try:
            for window in self.driver.window_handles:
                self.driver.switch_to.window(window)
                if ":VISIBLE" in self.driver.title and "PAUSE" not in self.driver.title:
                    # CrawlerLogger.info(f"{self.driver.title} is top window")
                    return
            for window in self.driver.window_handles:
                self.driver.switch_to.window(window)
                if ":" not in self.driver.title and self.driver.title != '' and 'INVISIBLE' not in self.driver.title:
                    return
        except Exception as e:
            CrawlerLogger.debug(f"[-] Error finding top window: {e}")
            pass

    def switch_to_webview_context(self):
        '''
        切换成当前driver的context成webview
        '''
        time.sleep(5)
        if self.driver.current_context != 'WEBVIEW_com.tencent.mm:appbrand' + conf.CLIENT_SIDE_CONF.WX_APP_BRAND:
            if ('WEBVIEW_com.tencent.mm:appbrand' + conf.CLIENT_SIDE_CONF.WX_APP_BRAND) in self.driver.contexts:
                CrawlerLogger.debug(f"{self.driver.current_context}")

                self.driver.switch_to.context('WEBVIEW_com.tencent.mm:appbrand' + conf.CLIENT_SIDE_CONF.WX_APP_BRAND)
                CrawlerLogger.debug(f"{self.driver.current_context}")

            else:
                CrawlerLogger.error("[-] Error finding webview context, maybe is a game")
                exit()
        max_retries = 10
        retry_count = 0
        while retry_count < max_retries:
            try:
                self.find_top_window()
                break
            except Exception as e:
                CrawlerLogger.debug(f"[-] Error finding top window: {e}")
                retry_count += 1
                time.sleep(0.5)  # 等待 500 毫秒
        else:
            CrawlerLogger.debug("[-] Failed to find top window after max retries")

    def get_url( self ) -> (str, str):
        '''
        获取当前页面的url, 统一格式 [path, query]
        :return:
        '''
        if self.driver.current_context == "NATIVE_APP" :
            self.switch_to_webview_context()
        try :
            self.find_top_window()
            url = self.get_wechat_url()
            # 动态更新用于跨进程通信
            conf.MINIAPP_CONF.pageUrl.value = url
            path = url.split("?")[0].split(".")[0]
            query = url.split("?")[1] if len(url.split("?")) > 1 else ""
            return path, query
        except Exception as e :
            CrawlerLogger.debug(f"Error getting driver title: {e}")
            return "pages/index/index", ""

    def get_all_pages( self, file_path ):
        with open(file_path, 'r', encoding='utf-8') as f :
            json_data = json.load(f)
        all_page = json_data['pages']
        return all_page

    def enter_mini_program( self, applets_name):
        '''
        进入小程序
        :param applets_name:
        :return:
        '''
        self.detach_from_app()
        self.switch_to_native_app_context()
        open_app("com.tencent.mm", "com.tencent.mm.ui.LauncherUI", conf.CLIENT_SIDE_CONF.DEVICE_ID)
        self.attach_to_app('微信', 'wechat')
        time.sleep(5)
        
        # # 等待页面加载完毕
        # WebDriverWait(self.driver, 10).until(
        #     EC.presence_of_element_located((AppiumBy.XPATH, "//*[@text='通讯录']"))
        # )
        # # 进入小程序
        # CrawlerLogger.info(f"[*] 进入小程序...")

        # PageHandler.safe_click_element(self.driver, AppiumBy.ACCESSIBILITY_ID, "搜索")
        # WebDriverWait(self.driver, 10).until(
        #     EC.presence_of_element_located((AppiumBy.XPATH, "//*[@text='搜索指定内容']"))
        # )
        # input = PageHandler.safe_find_element(self.driver, AppiumBy.ID, "com.tencent.mm:id/bxz")
        # input.send_keys(applets_name)
        # PageHandler.safe_click_element(self.driver, AppiumBy.ID, "com.tencent.mm:id/jkc")
        # time.sleep(2)

        # # TODO how click microprogram button in tool webview
        # action = TouchAction(self.driver)
        # time.sleep(4)
        # action.tap(None, 0.750 * conf.CLIENT_SIDE_CONF.PHONE_WIDTH, 0.143 * conf.CLIENT_SIDE_CONF.PHONE_HEIGHT).perform()
        # time.sleep(4)
        # action.tap(None, 0.138 * conf.CLIENT_SIDE_CONF.PHONE_WIDTH, 0.328 * conf.CLIENT_SIDE_CONF.PHONE_HEIGHT).perform()
        # time.sleep(4)

        # frida 方式开启小程序
        self.startMiniApp(conf.MINIAPP_CONF.applets_name.value)

        self.driver.quit()
        options = UiAutomator2Options().load_capabilities(self.CAPS)
        self.driver = webdriver.Remote("http://localhost:4723", options=options)
        # self.driver = webdriver.Remote('http://host.docker.internal:4723', options=options)
        self.driver.implicitly_wait(5)
        if PageHandler.safe_find_element(self.driver, AppiumBy.XPATH,
                                         f"//*[contains(@text, '网页由 mp.weixin.qq.com 提供')]", timeout=1) :
            CrawlerLogger.error("[*] 版本问题，结束crawler")
            with open(f'./experiment/wx_files/large_scale/applet_info.json', 'r', encoding='utf-8') as f :
                applet_info = json.load(f)
            applet_info[self.applets_name] = "版本问题"
            with open(f'./experiment/wx_files/large_scale/applet_info.json', 'w', encoding='utf-8') as f :
                json.dump(applet_info, f, ensure_ascii=False, indent=4)
            exit()

        CrawlerLogger.info(f"[*] 进入小程序成功...")
        time.sleep(5)
        # if self.get_wechat_appid():
        #     CrawlerLogger.info(f"[*] 进入小程序成功...")
        # else:
        #     CrawlerLogger.error(f"[-] 进入小程序失败...")
        #     self.enter_mini_program(applets_name)


    def get_app_json( self, applets_name ):
        # action = TouchAction(self.driver)
        # action.tap(None, 1144, 254).perform()
        # WebDriverWait(self.driver, 10).until(
        #     EC.presence_of_element_located((AppiumBy.XPATH, "//*[@text='发送给朋友']"))
        # )
        # action.tap(None, 308, 1579).perform()
        # WebDriverWait(self.driver, 10).until(
        #     EC.presence_of_element_located((AppiumBy.XPATH, "//*[@text='更多资料']"))
        # )
        # PageHandler.safe_click_element(self.driver, AppiumBy.XPATH, "//*[@text='更多资料']")
        # self.driver.switch_to.context('WEBVIEW_com.tencent.mm:toolsmp')
        #
        # # 等待页面加载完毕
        # WebDriverWait(self.driver, 10).until(
        #     EC.presence_of_element_located((AppiumBy.CSS_SELECTOR, 'body'))
        # )
        #
        # # 获取页面源码
        # page_source = self.driver.page_source
        #
        # # 使用正则表达式来提取AppID
        # match = re.search(r'AppID<\/h3>\s*<div class="info_item_desc">(wx[0-9a-z]+)<\/div>', page_source)
        try:
            # self.appid = self.get_wechat_appid()
            self.appid = conf.MINIAPP_CONF.applets_name.value
            with open(f'./experiment/wx_files/large_scale/applet_info.json', 'r', encoding='utf-8') as f :
                applet_info = json.load(f)
            applet_info[self.applets_name] = self.appid
            with open(f'./experiment/wx_files/large_scale/applet_info.json', 'w', encoding='utf-8') as f :
                json.dump(applet_info, f, ensure_ascii=False, indent=4)

            if not os.path.exists('./experiment/wx_files/pkgs_handle/' + self.applets_name) :
                os.mkdir('./experiment/wx_files/pkgs_handle/' + self.applets_name)
            copy_Pkg_Down_wx(self.user_hash, self.appid, './experiment/wx_files/pkgs_handle/' + self.applets_name, conf.CLIENT_SIDE_CONF.DEVICE_ID)
            unpack_and_copy_json_wx(applets_name, './experiment/wx_files/pkgs_handle/' + self.applets_name, conf.WX_UNPACK_PATH, './experiment/app_json/wx')
        except Exception as e:
            CrawlerLogger.error(f"[-] {traceback.format_exc()}")
            CrawlerLogger.debug(f"[-] {self.applets_name}未找到AppID")
            exit()

    def go_to_target_url ( self, target_path: str, query: str = "" ):
        '''
        # driver_handler.driver.execute_script('wx.reLaunch({url: "/pages/policy/index"})')
        :param target_path:
        :param query:
        :return:
        '''
        if self.driver.current_context == "NATIVE_APP":
            self.switch_to_webview_context()

        # Construct the URL
        base_url = target_path.split(":")[1].split(".")[0] if ":" in target_path else target_path
        if query:
            url = f"/{base_url}?{query}"
        else:
            url = f"/{base_url}"
        js_script = f'wx.reLaunch({{url: "{url}"}})'

        CrawlerLogger.info(f"[*] js go to {target_path}...")
        self.driver.execute_script(js_script)
        time.sleep(2)
        self.find_top_window()
